"""
$Id: engine.py 2 2007-03-04 19:08:44Z sidnei.da.silva $
"""

from evolve.signals import *
from evolve.choice import choice, shuffle
from evolve.generation import Generation
from evolve.genome import Genome

class Evolve(object):

    def __init__(self, fitness, pool, crossover, mutate, stop, notify):
        """Initialize evolution engine.

        fitness - A callable to calculate the fitness value for a genome
        pool - First generation pool of genomes
        crossover - A callable to crossover two genomes and create a new one
        mutate - A callable to mutate a genome
        stop - A callable to check for the stop condition
        notify - A callable that will receive events
        """
        self.fitness = fitness
        self.current = Generation(*pool)
        self.crossover = crossover
        self.mutate = mutate
        self.stop = stop
        self.notify = notify
        self.immigrate = False
        self.elitize = True

    def step(self):
        """Run evolutionary algorithm through one generation,
        returning the next generation and setting it as the current
        generation.
        """
        generation = self.current
        fit = []
        for genome in generation:
            fit.append((self.fitness(genome), genome))
        next = []
        for x, y in self.pairs(fit):
            for child in self.crossover(x, y):
                child = self.mutate(child)
                if self.elitize:
                    if self.fitness(child) < self.fitness(x):
                        child = x
                next.append(child)
        self.current = Generation(*next)
        self.notify(STEP, generation, self.current)
        return self.current

    def pairs(self, fit):
        # Order by fitness
        fit.sort()
        # Reverse so it's in descending order
        fit.reverse()
        n_pairs = len(fit)
        best = (n_pairs/4) * 3
        # Select the 75% best as the 'best breed' genome
        chosen = fit[:best]
        while len(chosen) < (n_pairs):
            # Choose one genome to duplicate until we reach n_pairs - 1
            dup = choice(chosen)
            chosen.append(dup)
        # Extract only the genome, dropping the fit value
        fit = [y for x, y in chosen]
        if self.immigrate:
            # Leak in a 'random immigrant' to avoid reaching a 'local
            # optima' too fast
            immigrant = list(choice(fit))
            shuffle(immigrant)
            fit.pop()
            fit.append(Genome(*immigrant))
        pairs = []
        for a in fit:
            for i in fit:
                b = choice(fit)
                if a != b:
                    break
            pairs.append((a, b))
        return pairs

    def run(self):
        self.notify(START)
        while not self.stop(self.current):
            self.step()
        self.notify(FINISH, (self.current,))
